import netCDF4
import numpy as np
import glob #用于查找符合规则的路径名
import os

month = "10"
fpath = "path/to/file" + month + "/OUT_D12_WRF"
fname_pattern = os.path.join(fpath, "wrfout_d01_*")
fname = glob.glob(fname_pattern)
nfile = len(fname)
#用来查找符合我们指定的wrfout文件

ncol = 232
nrow = 182
nlay = 23

vars_list = ["tc", "rh", "p", "VBS01", "VBS02", "VBS03", "VBS04", "VBS05", "VBS06", "VBS07", "VBS08", "VBS09", "VBS10", "alk3", "alk4", "alk5", "aro1", "aro2", "pcg7_f_c",
            "terp", "sesq", "isoprene", "ho",
            "DVBS01DT", "DVBS02DT", "DVBS03DT", "DVBS04DT", "DVBS05DT", "DVBS06DT", "DVBS07DT", "DVBS08DT", "DVBS09DT", "DVBS10DT"]
#变量列表感觉可以优化一下，变量列表用的是张华师兄的NCL脚本的变量，一些其他文献我还需要再查查

varnames = vars_list 
nvar = len(vars_list)

nsize = nrow * ncol * nlay * nfile
values = np.zeros((nsize, nvar), dtype=np.float32) 

for f_index, file_name in enumerate(fname):
    print(f"Processing file: {file_name}")
    fin = netCDF4.Dataset(file_name, "r")
    for v_index, var_name in enumerate(vars_list):
        try:

            temp0 = fin.variables[var_name][:] # Read all time steps and dimensions

            if len(temp0.shape) > 3: # if time, z, y, x
                temp0_single_timestep = temp0[0, :, :, :] # Take the first (and likely only) timestep
            elif len(temp0.shape) > 2: # if time, y, x
                temp0_single_timestep = temp0[0, :, :] # Take the first (and likely only) timestep
            else:
                temp0_single_timestep = temp0 # or just assign it if it's already 3D without time dim

            temp2 = temp0_single_timestep.flatten() # Flatten to 1D array
            values[f_index * nrow * ncol * nlay:(f_index + 1) * nrow * ncol * nlay, v_index] = temp2
        except KeyError:
            print(f"Variable '{var_name}' not found in file '{file_name}'. Skipping.")
            # If variable is not found, fill with NaN or zeros, or handle as needed.
            values[f_index * nrow * ncol * nlay:(f_index + 1) * nrow * ncol * nlay, v_index] = np.nan # Fill with NaN for missing variables

    fin.close()

output_file = "./data/2D-VBS_" + month + ".txt"
np.savetxt(output_file, values, fmt='%15.4e') #这里借鉴了张华师兄的NCL脚本，NCL对精度控制好像和Python不一样，这个我暂时不太清楚

print(f"Data saved to {output_file}")